#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>

#pragma pack(1)
typedef struct {
    uint8_t magic1;        // 0xAA
    uint8_t magic2;        // 0x55
    uint8_t fw_count;      // 固定1
    uint8_t major_version;
    uint8_t minor_version;
    uint8_t reserved[11];  // 全0
    uint32_t fw_length;    // 固件长度（小端）
    uint16_t checksum;     // 校验和（小端）
    char fw_type[10];      // 固件类型字符串
} OTAHeader;

// 转换.cly文件为bin格式
int cly_to_bin(const char* input, const char* output) {
    FILE* fin = fopen(input, "r");
    FILE* fout = fopen(output, "wb");
    if (!fin || !fout) {
        perror("File open failed");
        return -1;
    }

    // 读取首行4字节大端长度
    char len_str[9] = {0};
    fread(len_str, 1, 8, fin);
    uint32_t data_len = strtoul(len_str, NULL, 16);
    data_len = (data_len >> 24) |          // 大端转小端
               ((data_len >> 8) & 0xFF00) |
               ((data_len << 8) & 0xFF0000) |
               (data_len << 24);

    char line[129]; // 每行128字符（64字节）
    uint32_t total_bytes = 0;
    
    while (fgets(line, sizeof(line), fin)) {
        // 去除换行符
        line[strcspn(line, "\r\n")] = 0;
        
        // 逐字节转换HEX字符串
        for (int i = 0; i < strlen(line); i += 2) {
            uint8_t byte = 0;
            if (sscanf(line + i, "%2hhx", &byte) != 1) {
                fprintf(stderr, "Invalid HEX format at line %d\n", __LINE__);
                fclose(fin);
                fclose(fout);
                return -1;
            }
            fwrite(&byte, 1, 1, fout);
            total_bytes++;
            
            if (total_bytes >= data_len) break;
        }
    }

    // 验证数据长度
    if (total_bytes != data_len) {
        fprintf(stderr, "Data length mismatch! Header:%u Actual:%u\n", 
               data_len, total_bytes);
        fclose(fin);
        fclose(fout);
        return -1;
    }

    fclose(fin);
    fclose(fout);
    return 0;
}

// 计算校验和（累加和）
uint16_t calculate_checksum(FILE* fp) {
    uint16_t sum = 0;
    uint8_t byte;
    while (fread(&byte, 1, 1, fp) == 1) {
        sum += byte;
    }
    rewind(fp);
    return sum;
}

// 打包功能实现
int pack_firmware(const char* filename, const char* type, const char* version) {
    /* 1. 文件格式转换处理 */
    char tmpfile[] = "temp.bin";
    const char* ext = strrchr(filename, '.');

    if (ext) {
        if (strcmp(ext, ".bin") == 0) {
            // 直接复制bin文件到临时文件
            FILE* src = fopen(filename, "rb");
            FILE* dst = fopen(tmpfile, "wb");
            uint8_t buf[4096];
            size_t bytes;
            while ((bytes = fread(buf, 1, sizeof(buf), src)) > 0) {
                fwrite(buf, 1, bytes, dst);
            }
            fclose(src);
            fclose(dst);
        }
        else if (strcmp(ext, ".hex") == 0) {
            // 调用hex2bin工具转换
            char cmd[256];
            snprintf(cmd, sizeof(cmd), "hex2bin.exe %s %s", filename, tmpfile);
            if (system(cmd) != 0) {
                fprintf(stderr, "Hex conversion failed\n");
                exit(EXIT_FAILURE);
            }
        }
        else if (strcmp(ext, ".cly") == 0) {
            // 调用cly转换函数
            if (cly_to_bin(filename, tmpfile) != 0) {
                fprintf(stderr, "CLY file conversion failed\n");
                exit(EXIT_FAILURE);
            }
        }
        else {
            fprintf(stderr, "Unsupported file format: %s\n", ext);
            exit(EXIT_FAILURE);
        }
    } else {
        fprintf(stderr, "Missing file extension\n");
        exit(EXIT_FAILURE);
    }

    /* 2. 准备OTA文件头 */
    OTAHeader header = {
        .magic1 = 0xAA,
        .magic2 = 0x55,
        .fw_count = 1,
        .major_version = (uint8_t)atoi(strtok(version, ".")),
        .minor_version = (uint8_t)atoi(strtok(NULL, ".")),
    };
    strncpy(header.fw_type, type, sizeof(header.fw_type));

    /* 3. 写入文件 */
    // 生成输出文件名
    char output_file[256];
    strncpy(output_file, filename, sizeof(output_file));
    char* oext = strrchr(output_file, '.');  // 查找文件扩展名
    if (oext) *oext = '\0';                   // 移除原扩展名
    strcat(output_file, "_ota.bin");

    // 获取固件数据长度和校验和
    FILE* tmp_fp = fopen("temp.bin", "rb"); // 假设转换后的临时文件
    fseek(tmp_fp, 0, SEEK_END);
    header.fw_length = ftell(tmp_fp);      // 获取固件长度
    rewind(tmp_fp);
    header.checksum = calculate_checksum(tmp_fp);

    // 写入OTA文件
    FILE* out_fp = fopen(output_file, "wb");
    fwrite(&header, sizeof(OTAHeader), 1, out_fp);  // 写文件头
    uint8_t buffer[4096];
    size_t bytes_read;
    while ((bytes_read = fread(buffer, 1, sizeof(buffer), tmp_fp)) > 0) {
        fwrite(buffer, 1, bytes_read, out_fp);      // 写固件数据
    }

    // 清理资源
    fclose(tmp_fp);
    fclose(out_fp);
    remove("temp.bin");  // 删除临时文件
    return 0;

}

// 主函数（命令行解析）
int main(int argc, char* argv[]) {
    int opt;
    OTAHeader header;

    while ((opt = getopt(argc, argv, "hp:u:v:")) != -1) {
        switch (opt) {
        case 'p': // 打包
            if (argc != optind + 2) {
                fprintf(stderr, "Usage: %s -p [file] [type] [version]\n", argv[0]);
                exit(EXIT_FAILURE);
            }
            pack_firmware(optarg, argv[optind], argv[optind+1]);
            break;
        case 'u': // 解包
            if (argc != optind) {
                fprintf(stderr, "Usage: %s -u [file]\n", argv[0]);
                exit(EXIT_FAILURE);
            }
            FILE* in_fp = fopen(optarg, "rb");
            if (!in_fp) {
                perror("Open OTA file failed");
                exit(EXIT_FAILURE);
            }

            if (fread(&header, sizeof(OTAHeader), 1, in_fp) != 1) {
                fprintf(stderr, "Invalid OTA file format\n");
                fclose(in_fp);
                exit(EXIT_FAILURE);
            }

            // 验证文件头有效性
            if (header.magic1 != 0xAA || header.magic2 != 0x55 || header.fw_count != 1) {
                fprintf(stderr, "Invalid OTA file header\n");
                fclose(in_fp);
                exit(EXIT_FAILURE);
            }

            // 生成输出文件名
            char output_file[256];
            strncpy(output_file, optarg, sizeof(output_file));
            char* ext = strrchr(output_file, '.');
            if (ext && strcmp(ext, ".bin") != 0) *ext = '\0';
            strcat(output_file, "_extracted.bin");

            // 写入解包文件
            FILE* out_fp = fopen(output_file, "wb");
            uint8_t buffer[4096];
            size_t total_read = 0;
            size_t bytes_read;
            
            while (total_read < header.fw_length && 
                  (bytes_read = fread(buffer, 1, sizeof(buffer), in_fp)) > 0) {
                fwrite(buffer, 1, bytes_read, out_fp);
                total_read += bytes_read;
            }

            // 验证文件完整性
            if (total_read != header.fw_length) {
                fprintf(stderr, "File incomplete! Expected:%u Actual:%zu\n",
                       header.fw_length, total_read);
                remove(output_file);
                exit(EXIT_FAILURE);
            }

            printf("Successfully unpackaged to: %s\n", output_file);
            fclose(in_fp);
            fclose(out_fp);
            break;

        case 'v': // 查看信息
            if (argc != optind) {
                fprintf(stderr, "Usage: %s -v [file]\n", argv[0]);
                exit(EXIT_FAILURE);
            }
            FILE* ota_fp = fopen(optarg, "rb");
            if (!ota_fp) {
                perror("Open OTA file failed");
                exit(EXIT_FAILURE);
            }

            if (fread(&header, sizeof(OTAHeader), 1, ota_fp) != 1) {
                fprintf(stderr, "Invalid OTA file format\n");
                fclose(ota_fp);
                exit(EXIT_FAILURE);
            }

            // 验证文件头标识
            if (header.magic1 != 0xAA || header.magic2 != 0x55) {
                fprintf(stderr, "Invalid magic number\n");
                fclose(ota_fp);
                exit(EXIT_FAILURE);
            }
            // 验证checksum是否正确
            fseek(ota_fp, sizeof(OTAHeader), SEEK_SET); // 跳过文件头
            uint16_t actual_checksum = 0;
            uint8_t byte;
            while (fread(&byte, 1, 1, ota_fp) == 1) {
                actual_checksum += byte;
            }
            
            if (actual_checksum != header.checksum) {
                fprintf(stderr, "Checksum mismatch! File:0x%04X, Calculated:0x%04X\n",
                       header.checksum, actual_checksum);
                fclose(ota_fp);
                exit(EXIT_FAILURE);
            }

            printf("OTA File Information:\n"
                   "Firmware Type:  %.10s\n"
                   "Version:        %d.%d\n"
                   "Firmware Size:  %u bytes\n"
                   "Checksum:       0x%04X\n",
                   header.fw_type,
                   header.major_version,
                   header.minor_version,
                   header.fw_length,
                   header.checksum);

            fclose(ota_fp);
            break;

        case 'h': // 帮助信息
            printf("OTA Packer Tool v1.0\n"
                   "Usage: %s [options]\n"
                   "Options:\n"
                   "  -p <file> <type> <version>  Package firmware\n"
                   "  -u <file>                  Unpackage OTA file\n"
                   "  -v <file>                  Verify OTA file\n"
                   "Supported formats: .bin, .hex, .cly\n", argv[0]);
            break;
        default:
            fprintf(stderr, "Invalid option\n");
            exit(EXIT_FAILURE);
        }
    }
    return 0;
}
